package com.three.db.schema;

import com.three.db.config.DataSourceConfig;
import org.apache.commons.collections4.CollectionUtils;
import org.hibernate.boot.Metadata;
import org.hibernate.boot.model.naming.Identifier;
import org.hibernate.boot.model.relational.*;
import org.hibernate.boot.model.relational.AuxiliaryDatabaseObject;
import org.hibernate.boot.spi.MetadataImplementor;
import org.hibernate.cfg.AvailableSettings;
import org.hibernate.dialect.Dialect;
import org.hibernate.engine.config.spi.ConfigurationService;
import org.hibernate.engine.config.spi.StandardConverters;
import org.hibernate.engine.jdbc.env.spi.JdbcEnvironment;
import org.hibernate.internal.util.StringHelper;
import org.hibernate.mapping.*;
import org.hibernate.tool.hbm2ddl.UniqueConstraintSchemaUpdateStrategy;
import org.hibernate.tool.schema.extract.spi.*;
import org.hibernate.tool.schema.spi.Exporter;
import org.hibernate.tool.schema.spi.SchemaManagementException;
import org.hibernate.tool.schema.spi.SchemaMigrator;
import org.hibernate.tool.schema.spi.Target;

import java.util.*;
import java.util.List;
import java.util.Set;

/**
 * Created by ziqingwang on 2017/6/20 0020.
 */
public class SupportShardingSchemaMigratorImpl implements SchemaMigrator {
    @Override
    public void doMigration(
            Metadata metadata,
            DatabaseInformation existingDatabase,
            boolean createNamespaces,
            List<Target> targets) throws SchemaManagementException {

        for ( Target target : targets ) {
            target.prepare();
        }

        doMigrationToTargets( metadata, existingDatabase, createNamespaces, targets );

        for ( Target target : targets ) {
            target.release();
        }
    }


    protected void doMigrationToTargets(
            Metadata metadata,
            DatabaseInformation existingDatabase,
            boolean createNamespaces,
            List<Target> targets) {
        final Set<String> exportIdentifiers = new HashSet<String>( 50 );

        final Database database = metadata.getDatabase();
        final JdbcEnvironment jdbcEnvironment = database.getJdbcEnvironment();
        final Dialect dialect = jdbcEnvironment.getDialect();

        // Drop all AuxiliaryDatabaseObjects
        for ( AuxiliaryDatabaseObject auxiliaryDatabaseObject : database.getAuxiliaryDatabaseObjects() ) {
            if ( !auxiliaryDatabaseObject.appliesToDialect( dialect ) ) {
                continue;
            }

            applySqlStrings(
                    dialect.getAuxiliaryDatabaseObjectExporter().getSqlDropStrings( auxiliaryDatabaseObject, metadata ),
                    targets,
                    true
            );
        }

        // Create before-table AuxiliaryDatabaseObjects
        for ( AuxiliaryDatabaseObject auxiliaryDatabaseObject : database.getAuxiliaryDatabaseObjects() ) {
            if ( auxiliaryDatabaseObject.beforeTablesOnCreation() ) {
                continue;
            }
            if ( !auxiliaryDatabaseObject.appliesToDialect( dialect ) ) {
                continue;
            }

            applySqlStrings(
                    auxiliaryDatabaseObject.sqlCreateStrings( jdbcEnvironment.getDialect() ),
                    targets,
                    true
            );
        }

        boolean tryToCreateCatalogs = false;
        boolean tryToCreateSchemas = false;

        if ( createNamespaces ) {
            if ( database.getJdbcEnvironment().getDialect().canCreateSchema() ) {
                tryToCreateSchemas = true;
            }
            if ( database.getJdbcEnvironment().getDialect().canCreateCatalog() ) {
                tryToCreateCatalogs = true;
            }
        }

        Set<Identifier> exportedCatalogs = new HashSet<Identifier>();
        for ( Namespace namespace : database.getNamespaces() ) {
            if ( tryToCreateCatalogs || tryToCreateSchemas ) {
                if ( tryToCreateCatalogs ) {
                    final Identifier catalogLogicalName = namespace.getName().getCatalog();
                    final Identifier catalogPhysicalName = namespace.getPhysicalName().getCatalog();

                    if ( catalogPhysicalName != null && !exportedCatalogs.contains( catalogLogicalName ) && !existingDatabase
                            .catalogExists( catalogLogicalName ) ) {
                        applySqlStrings(
                                database.getJdbcEnvironment().getDialect().getCreateCatalogCommand(
                                        catalogPhysicalName.render(
                                                database.getJdbcEnvironment().getDialect()
                                        )
                                ),
                                targets,
                                false
                        );
                        exportedCatalogs.add( catalogLogicalName );
                    }
                }

                if ( tryToCreateSchemas
                        && namespace.getPhysicalName().getSchema() != null
                        && !existingDatabase.schemaExists( namespace.getName() ) ) {
                    applySqlStrings(
                            database.getJdbcEnvironment().getDialect().getCreateSchemaCommand(
                                    namespace.getPhysicalName()
                                            .getSchema()
                                            .render( database.getJdbcEnvironment().getDialect() )
                            ),
                            targets,
                            false
                    );
                }
            }

            for ( Table table : namespace.getTables() ) {
                if ( !table.isPhysicalTable() ) {
                    continue;
                }
                TableSchema schema= DataSourceConfig.SHARDING_TABLE_NAME_MAP.get(table.getName());
                if(Objects.nonNull(schema)) {//分片数据表
                    Set<String> dbSet = schema.getShardingDataBase();
                    Set<String> tbSet = schema.getShardingTable();
                    if(CollectionUtils.isNotEmpty(dbSet)&&CollectionUtils.isNotEmpty(tbSet)){
                        dbSet.stream().forEach(db->{
                            tbSet.stream().forEach(tb->{
                                table.setName(tb);
                                table.setCatalog(db);
                                checkExportIdentifier( table, exportIdentifiers );
                                final TableInformation tableInformation = existingDatabase.getTableInformation( table.getQualifiedTableName() );
                                if ( tableInformation != null && !tableInformation.isPhysicalTable() ) {
                                    return;
                                }
                                if ( tableInformation == null ) {
                                    createTable( table, metadata, targets );
                                }
                                else {
                                    migrateTable( table, tableInformation, targets, metadata );
                                }
                            });
                        });
                    }
                }else{
                    checkExportIdentifier( table, exportIdentifiers );
                    final TableInformation tableInformation = existingDatabase.getTableInformation( table.getQualifiedTableName() );
                    if ( tableInformation != null && !tableInformation.isPhysicalTable() ) {
                        continue;
                    }
                    if ( tableInformation == null ) {
                        createTable( table, metadata, targets );
                    }
                    else {
                        migrateTable( table, tableInformation, targets, metadata );
                    }
                }
            }

            for ( Table table : namespace.getTables() ) {
                if ( !table.isPhysicalTable() ) {
                    continue;
                }
                TableSchema schema= DataSourceConfig.SHARDING_TABLE_NAME_MAP.get(table.getName());
                if(Objects.nonNull(schema)) {//分片数据表
                    Set<String> dbSet = schema.getShardingDataBase();
                    Set<String> tbSet = schema.getShardingTable();
                    if(CollectionUtils.isNotEmpty(dbSet)&&CollectionUtils.isNotEmpty(tbSet)){
                        dbSet.stream().forEach(db->{
                            tbSet.stream().forEach(tb->{
                                table.setName(tb);
                                table.setCatalog(db);
                                final TableInformation tableInformation = existingDatabase.getTableInformation( table.getQualifiedTableName() );
                                if ( tableInformation != null && !tableInformation.isPhysicalTable() ) {
                                    return;
                                }
                                applyIndexes( table, tableInformation, metadata, targets );
                                applyUniqueKeys( table, tableInformation, metadata, targets );
                            });
                        });
                    }
                }else{
                    final TableInformation tableInformation = existingDatabase.getTableInformation( table.getQualifiedTableName() );
                    if ( tableInformation != null && !tableInformation.isPhysicalTable() ) {
                        continue;
                    }
                    applyIndexes( table, tableInformation, metadata, targets );
                    applyUniqueKeys( table, tableInformation, metadata, targets );
                }
            }

            for ( Sequence sequence : namespace.getSequences() ) {
                checkExportIdentifier( sequence, exportIdentifiers );
                final SequenceInformation sequenceInformation = existingDatabase.getSequenceInformation( sequence.getName() );
                if ( sequenceInformation != null ) {
                    // nothing we really can do...
                    continue;
                }

                applySqlStrings(
                        database.getJdbcEnvironment().getDialect().getSequenceExporter().getSqlCreateStrings(
                                sequence,
                                metadata
                        ),
                        targets,
                        false
                );
            }
        }

        //NOTE : Foreign keys must be created *after* all tables of all namespaces for cross namespace fks. see HHH-10420
        for ( Namespace namespace : database.getNamespaces() ) {
            for ( Table table : namespace.getTables() ) {
                final TableInformation tableInformation = existingDatabase.getTableInformation( table.getQualifiedTableName() );
                if ( tableInformation != null && !tableInformation.isPhysicalTable() ) {
                    continue;
                }
                applyForeignKeys( table, tableInformation, metadata, targets );
            }
        }

        // Create after-table AuxiliaryDatabaseObjects
        for ( AuxiliaryDatabaseObject auxiliaryDatabaseObject : database.getAuxiliaryDatabaseObjects() ) {
            if ( !auxiliaryDatabaseObject.beforeTablesOnCreation() ) {
                continue;
            }
            if ( !auxiliaryDatabaseObject.appliesToDialect( dialect ) ) {
                continue;
            }

            applySqlStrings(
                    auxiliaryDatabaseObject.sqlCreateStrings( jdbcEnvironment.getDialect() ),
                    targets,
                    true
            );
        }
    }

    private void createTable(Table table, Metadata metadata, List<Target> targets) {
        applySqlStrings(
                metadata.getDatabase().getDialect().getTableExporter().getSqlCreateStrings( table, metadata ),
                targets,
                false
        );
    }

    private void migrateTable(
            Table table,
            TableInformation tableInformation,
            List<Target> targets,
            Metadata metadata) {
        final Database database = metadata.getDatabase();
        final JdbcEnvironment jdbcEnvironment = database.getJdbcEnvironment();
        final Dialect dialect = jdbcEnvironment.getDialect();

        //noinspection unchecked
        applySqlStrings(
                table.sqlAlterStrings(
                        dialect,
                        metadata,
                        tableInformation,
                        getDefaultCatalogName( database ),
                        getDefaultSchemaName( database )
                ),
                targets,
                false
        );
    }

    private void applyIndexes(Table table, TableInformation tableInformation, Metadata metadata, List<Target> targets) {
        final Exporter<Index> exporter = metadata.getDatabase().getJdbcEnvironment().getDialect().getIndexExporter();

        final Iterator<Index> indexItr = table.getIndexIterator();
        while ( indexItr.hasNext() ) {
            final Index index = indexItr.next();
            if ( StringHelper.isEmpty( index.getName() ) ) {
                continue;
            }

            if ( tableInformation != null ) {
                final IndexInformation existingIndex = findMatchingIndex( index, tableInformation );
                if ( existingIndex != null ) {
                    continue;
                }
            }

            applySqlStrings(
                    exporter.getSqlCreateStrings( index, metadata ),
                    targets,
                    false
            );
        }
    }

    private IndexInformation findMatchingIndex(Index index, TableInformation tableInformation) {
        return tableInformation.getIndex( Identifier.toIdentifier( index.getName() ) );
    }

    private UniqueConstraintSchemaUpdateStrategy uniqueConstraintStrategy;

    private void applyUniqueKeys(Table table, TableInformation tableInfo, Metadata metadata, List<Target> targets) {
        if ( uniqueConstraintStrategy == null ) {
            uniqueConstraintStrategy = determineUniqueConstraintSchemaUpdateStrategy( metadata );
        }

        if ( uniqueConstraintStrategy == UniqueConstraintSchemaUpdateStrategy.SKIP ) {
            return;
        }

        final Dialect dialect = metadata.getDatabase().getJdbcEnvironment().getDialect();
        final Exporter<Constraint> exporter = dialect.getUniqueKeyExporter();

        final Iterator ukItr = table.getUniqueKeyIterator();
        while ( ukItr.hasNext() ) {
            final UniqueKey uniqueKey = (UniqueKey) ukItr.next();
            // Skip if index already exists. Most of the time, this
            // won't work since most Dialects use Constraints. However,
            // keep it for the few that do use Indexes.
            if ( tableInfo != null && StringHelper.isNotEmpty( uniqueKey.getName() ) ) {
                final IndexInformation indexInfo = tableInfo.getIndex( Identifier.toIdentifier( uniqueKey.getName() ) );
                if ( indexInfo != null ) {
                    continue;
                }
            }

            if ( uniqueConstraintStrategy == UniqueConstraintSchemaUpdateStrategy.DROP_RECREATE_QUIETLY ) {
                applySqlStrings(
                        exporter.getSqlDropStrings( uniqueKey, metadata ),
                        targets,
                        true
                );
            }

            applySqlStrings(
                    exporter.getSqlCreateStrings( uniqueKey, metadata ),
                    targets,
                    true
            );
        }
    }

    private UniqueConstraintSchemaUpdateStrategy determineUniqueConstraintSchemaUpdateStrategy(Metadata metadata) {
        final ConfigurationService cfgService = ((MetadataImplementor) metadata).getMetadataBuildingOptions()
                .getServiceRegistry()
                .getService( ConfigurationService.class );

        return UniqueConstraintSchemaUpdateStrategy.interpret(
                cfgService.getSetting(
                        AvailableSettings.UNIQUE_CONSTRAINT_SCHEMA_UPDATE_STRATEGY,
                        StandardConverters.STRING
                )
        );
    }

    private void applyForeignKeys(
            Table table,
            TableInformation tableInformation,
            Metadata metadata,
            List<Target> targets) {
        final Dialect dialect = metadata.getDatabase().getJdbcEnvironment().getDialect();
        if ( !dialect.hasAlterTable() ) {
            return;
        }

        final Exporter<ForeignKey> exporter = dialect.getForeignKeyExporter();

        @SuppressWarnings("unchecked")
        final Iterator<ForeignKey> fkItr = table.getForeignKeyIterator();
        while ( fkItr.hasNext() ) {
            final ForeignKey foreignKey = fkItr.next();
            if ( !foreignKey.isPhysicalConstraint() ) {
                continue;
            }

            if ( !foreignKey.isCreationEnabled() ) {
                continue;
            }

            if ( tableInformation != null ) {
                final ForeignKeyInformation existingForeignKey = findMatchingForeignKey( foreignKey, tableInformation );
                if ( existingForeignKey != null ) {
                    continue;
                }
            }

            // todo : shouldn't we just drop+recreate if FK exists?
            //		this follows the existing code from legacy SchemaUpdate which just skipped

            // in old SchemaUpdate code, this was the trigger to "create"
            applySqlStrings(
                    exporter.getSqlCreateStrings( foreignKey, metadata ),
                    targets,
                    false
            );
        }
    }

    private ForeignKeyInformation findMatchingForeignKey(ForeignKey foreignKey, TableInformation tableInformation) {
        if ( foreignKey.getName() == null ) {
            return null;
        }
        return tableInformation.getForeignKey( Identifier.toIdentifier( foreignKey.getName() ) );
    }

    private void checkExportIdentifier(Exportable exportable, Set<String> exportIdentifiers) {
        final String exportIdentifier = exportable.getExportIdentifier();
        if ( exportIdentifiers.contains( exportIdentifier ) ) {
            throw new SchemaManagementException(
                    String.format(
                            "Export identifier [%s] encountered more than once",
                            exportIdentifier
                    )
            );
        }
        exportIdentifiers.add( exportIdentifier );
    }

    private static void applySqlStrings(String[] sqlStrings, List<Target> targets, boolean quiet) {
        if ( sqlStrings == null ) {
            return;
        }

        for ( String sqlString : sqlStrings ) {
            applySqlString( sqlString, targets, quiet );
        }
    }

    private static void applySqlString(String sqlString, List<Target> targets, boolean quiet) {
        if ( sqlString == null ) {
            return;
        }

        for ( Target target : targets ) {
            try {
                target.accept( sqlString );
            }
            catch (SchemaManagementException e) {
                if ( !quiet ) {
                    throw e;
                }
                // otherwise ignore the exception
            }
        }
    }

    private static void applySqlStrings(Iterator<String> sqlStrings, List<Target> targets, boolean quiet) {
        if ( sqlStrings == null ) {
            return;
        }

        while ( sqlStrings.hasNext() ) {
            final String sqlString = sqlStrings.next();
            applySqlString( sqlString, targets, quiet );
        }
    }

    private String getDefaultCatalogName(Database database) {
        final Identifier identifier = database.getDefaultNamespace().getPhysicalName().getCatalog();
        return identifier == null ? null : identifier.render( database.getJdbcEnvironment().getDialect() );
    }

    private String getDefaultSchemaName(Database database) {
        final Identifier identifier = database.getDefaultNamespace().getPhysicalName().getSchema();
        return identifier == null ? null : identifier.render( database.getJdbcEnvironment().getDialect() );
    }
}
